#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import time
import subprocess
import csv
from pathlib import Path
from datetime import datetime
from Common.handle_qnx import QNXConnect

try:
    from tqdm import tqdm
    TQDM_AVAILABLE = True
except ImportError:
    TQDM_AVAILABLE = False
    print("提示: 未找到tqdm库，将不显示进度条。如需显示进度条，请安装tqdm: pip install tqdm")

# 任务状态定义
STATUS_PENDING = "未开始"
STATUS_RUNNING = "处理中"
STATUS_SUCCESS = "成功"
STATUS_FAILED = "失败"
STATUS_SKIPPED = "已跳过"

class QNXAutoTest:
    def __init__(self, material_folders, output_folder, 
                 push_path="/mnt/vendor/nfs/fota/cjm/0526/tmp", 
                 qnx_work_dir="/fota/cjm/0526", 
                 algorithm_cmd="./trackdemo --model ./ --input ./tmp/ --testcase 1 --width 1280 --height 960",
                 csv_file=None,
                 status_folder=None,
                 max_retries=3,
                 continue_from_csv=False,
                 device_serial=None):
        # 支持单个文件夹或文件夹列表
        if isinstance(material_folders, (str, Path)):
            self.material_folders = [Path(material_folders)]
        else:
            self.material_folders = [Path(folder) for folder in material_folders]
            
        self.output_folder = Path(output_folder)
        self.push_path = push_path
        self.qnx_work_dir = qnx_work_dir
        self.algorithm_cmd = algorithm_cmd
        self.max_retries = max_retries
        self.continue_from_csv = continue_from_csv
        self.device_serial = device_serial
        
        # 状态文件夹
        if status_folder:
            self.status_folder = Path(status_folder)
        else:
            # 默认使用项目根目录下的status文件夹
            self.status_folder = Path(__file__).parent / "status"
        
        # 确保状态文件夹存在
        os.makedirs(self.status_folder, exist_ok=True)
        
        # CSV文件路径
        if csv_file:
            self.csv_file = Path(csv_file)
        else:
            self.csv_file = self.status_folder / f"task_status_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
        
        # 任务状态字典
        self.tasks = {}
        
        # 确保输出文件夹存在
        os.makedirs(self.output_folder, exist_ok=True)
        
        # 为每个素材文件夹创建对应的输出子文件夹
        for folder in self.material_folders:
            folder_name = folder.name
            output_subfolder = self.output_folder / folder_name
            os.makedirs(output_subfolder, exist_ok=True)
        
        # 初始化QNX连接
        self.qnx_connector = QNXConnect(device_serial=device_serial)
        
        # 获取adb路径
        self.adb_path = self.qnx_connector.adb_path
        
    def run_adb_command(self, command):
        """运行adb命令并返回结果"""
        return self.qnx_connector.run_adb_command(command)
    
    def push_file(self, file_path):
        """推送文件到QNX系统"""
        print(f"正在推送文件: {file_path}")
        command = f"push {file_path} {self.push_path}"
        result = self.run_adb_command(command)
        return result.returncode == 0
    
    def pull_file(self, qnx_file_path, local_path):
        """从QNX系统拉取文件"""
        print(f"正在拉取文件: {qnx_file_path}")
        command = f"pull {qnx_file_path} {local_path}"
        result = self.run_adb_command(command)
        return result.returncode == 0
    
    def delete_qnx_file(self, qnx_file_path):
        """使用adb命令删除QNX系统中的文件"""
        print(f"正在删除QNX系统中的文件: {qnx_file_path}")
        command = f"shell rm -f {qnx_file_path}"
        result = self.run_adb_command(command)
        return result.returncode == 0
    
    def clean_qnx_files(self, file_name):
        """清理QNX系统中的素材文件和结果文件"""
        print(f"正在清理QNX系统中的文件: {file_name}")
        
        # 素材文件路径
        material_path = f"{self.push_path}/{file_name}"
        
        # 可能的结果文件路径
        result_path1 = f"{self.push_path}/{Path(file_name).stem}_0.txt"
        result_path2 = f"{self.push_path}/{Path(file_name).stem}.txt"
        
        # 删除素材文件
        self.delete_qnx_file(material_path)
        
        # 删除结果文件(尝试两种可能的命名格式)
        self.delete_qnx_file(result_path1)
        self.delete_qnx_file(result_path2)
        
        print(f"QNX系统文件清理完成: {file_name}")
    
    def run_algorithm(self, file_name):
        """连接到QNX系统并运行算法"""
        print(f"正在运行算法处理文件: {file_name}")
        
        try:
            # 连接到QNX系统
            qnx_session = self.qnx_connector.connect_to_qnx()
            
            # 切换到工作目录
            qnx_session.sendline(f"cd {self.qnx_work_dir}")
            qnx_session.expect(['#', '>'], timeout=5)
            
            # 运行算法
            print(f"正在执行命令: {self.algorithm_cmd}")
            qnx_session.sendline(self.algorithm_cmd)
            
            # 等待算法完成，查找 "[PASSED] TrackingTest" 标记
            timeout = 1800  # 30分钟超时
            success = False
            start_time = time.time()
            
            while time.time() - start_time < timeout:
                try:
                    index = qnx_session.expect([r'\[PASSED\] TrackingTest', '#', '>', 'Error', 'Failed'], timeout=10)
                    if index == 0:
                        print("算法执行成功！")
                        success = True
                        time.sleep(5)
                        break
                    elif index in [3, 4]:
                        print("算法执行出错！")
                        break
                except:
                    # 继续等待输出
                    pass
            
            if not success and time.time() - start_time >= timeout:
                print("算法执行超时！")
            
            # 退出QNX连接
            qnx_session.sendline("exit")
            qnx_session.expect(['#', '>'], timeout=5)
            qnx_session.sendline("exit")
            qnx_session.close()
            
            return success
        except Exception as e:
            print(f"运行算法时出错: {str(e)}")
            return False
    
    def init_csv_file(self, file_list):
        """初始化CSV文件，记录所有待处理的文件"""
        # 如果是从CSV继续执行，则读取CSV中的状态
        if self.continue_from_csv and self.csv_file.exists():
            self.load_tasks_from_csv()
            return
        
        # 创建新的CSV文件
        with open(self.csv_file, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['素材路径', '文件名', '状态', '重试次数', '开始时间', '完成时间', '错误信息'])
            
            for material_path, file_path in file_list:
                task_key = f"{material_path.name}/{file_path.name}"
                self.tasks[task_key] = {
                    'material_path': material_path,
                    'status': STATUS_PENDING,
                    'retries': 0,
                    'start_time': '',
                    'end_time': '',
                    'error': ''
                }
                writer.writerow([material_path.name, file_path.name, STATUS_PENDING, 0, '', '', ''])
        
        print(f"初始化任务CSV文件: {self.csv_file}")
    
    def load_tasks_from_csv(self):
        """从CSV文件加载任务状态"""
        if not self.csv_file.exists():
            print(f"CSV文件不存在: {self.csv_file}")
            return
        
        self.tasks = {}
        with open(self.csv_file, 'r', newline='') as f:
            reader = csv.DictReader(f)
            for row in reader:
                task_key = f"{row['素材路径']}/{row['文件名']}"
                self.tasks[task_key] = {
                    'material_path': row['素材路径'],
                    'status': row['状态'],
                    'retries': int(row['重试次数']),
                    'start_time': row['开始时间'],
                    'end_time': row['完成时间'],
                    'error': row['错误信息']
                }
        
        print(f"从CSV文件加载了 {len(self.tasks)} 个任务")
    
    def update_task_status(self, material_path_name, file_name, status, error=''):
        """更新任务状态并写入CSV文件"""
        task_key = f"{material_path_name}/{file_name}"
        if task_key not in self.tasks:
            return
        
        task = self.tasks[task_key]
        task['status'] = status
        
        # 记录开始时间
        if status == STATUS_RUNNING and not task['start_time']:
            task['start_time'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        # 记录结束时间和错误
        if status in [STATUS_SUCCESS, STATUS_FAILED, STATUS_SKIPPED]:
            task['end_time'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            if error:
                task['error'] = error
        
        # 更新CSV文件
        temp_csv = self.csv_file.with_suffix('.tmp')
        with open(temp_csv, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['素材路径', '文件名', '状态', '重试次数', '开始时间', '完成时间', '错误信息'])
            
            for key, info in self.tasks.items():
                material_folder, fname = key.split('/', 1)
                writer.writerow([
                    material_folder, 
                    fname, 
                    info['status'], 
                    info['retries'], 
                    info['start_time'], 
                    info['end_time'], 
                    info['error']
                ])
        
        # 替换原CSV文件
        os.replace(temp_csv, self.csv_file)
    
    def increment_retry_count(self, material_path_name, file_name):
        """增加文件的重试次数"""
        task_key = f"{material_path_name}/{file_name}"
        if task_key in self.tasks:
            self.tasks[task_key]['retries'] += 1
            return self.tasks[task_key]['retries']
        return 0
    
    def should_process_file(self, material_path_name, file_name):
        """判断是否应该处理该文件（基于任务状态和重试次数）"""
        if not self.continue_from_csv:
            return True
        
        task_key = f"{material_path_name}/{file_name}"
        if task_key not in self.tasks:
            return True
        
        task = self.tasks[task_key]
        
        # 如果任务已成功或已跳过，则不再处理
        if task['status'] in [STATUS_SUCCESS, STATUS_SKIPPED]:
            return False
        
        # 如果重试次数已达上限，则不再处理
        if task['retries'] >= self.max_retries:
            return False
        
        return True
    
    def clean_push_directory(self):
        """清空QNX系统上的推送目录"""
        print(f"正在清空QNX系统上的推送目录: {self.push_path}")
        # 删除目录下所有文件
        command = f"shell rm -rf {self.push_path}/*"
        result = self.run_adb_command(command)
        
        # 确保目录存在
        command = f"shell mkdir -p {self.push_path}"
        result = self.run_adb_command(command)
        
        print(f"QNX系统推送目录已清空: {self.push_path}")
        return result.returncode == 0
    
    def process_file(self, material_folder, file_path):
        """处理单个文件"""
        material_path_name = material_folder.name
        file_name = file_path.name
        print(f"\n处理文件: {material_path_name}/{file_name}")
        
        # 更新任务状态为处理中
        self.update_task_status(material_path_name, file_name, STATUS_RUNNING)
        
        # 判断是否已达到最大重试次数
        task_key = f"{material_path_name}/{file_name}"
        current_retries = self.tasks[task_key]['retries']
        if current_retries >= self.max_retries:
            print(f"文件 {material_path_name}/{file_name} 已达到最大重试次数 {self.max_retries}，跳过处理")
            self.update_task_status(material_path_name, file_name, STATUS_SKIPPED, f"达到最大重试次数 {self.max_retries}")
            return False
            
        # 清空推送目录
        self.clean_push_directory()
        
        # 推送文件到QNX系统
        if not self.push_file(file_path):
            error_msg = f"推送文件 {file_name} 失败"
            print(error_msg)
            self.increment_retry_count(material_path_name, file_name)
            self.update_task_status(material_path_name, file_name, STATUS_FAILED, error_msg)
            return False
        
        result_pulled = False
        
        # 运行算法
        if not self.run_algorithm(file_name):
            error_msg = f"运行算法处理 {file_name} 失败"
            print(error_msg)
            self.increment_retry_count(material_path_name, file_name)
            self.update_task_status(material_path_name, file_name, STATUS_FAILED, error_msg)
            # 即使算法执行失败，也要清理文件
            self.clean_qnx_files(file_name)
            return False
        
        # 创建素材文件夹对应的输出子文件夹路径
        output_subfolder = self.output_folder / material_path_name
        
        # 拉取结果文件（添加_0后缀）
        result_file_name = file_path.stem + "_0.txt"
        qnx_result_path = f"{self.push_path}/{result_file_name}"
        local_result_path = output_subfolder / result_file_name
        
        if not self.pull_file(qnx_result_path, local_result_path):
            print(f"拉取结果文件 {result_file_name} 失败，尝试其他可能的命名格式...")
            
            # 尝试没有_0后缀的命名方式
            result_file_name_alt = file_path.stem + ".txt"
            qnx_result_path_alt = f"{self.push_path}/{result_file_name_alt}"
            local_result_path_alt = output_subfolder / result_file_name_alt
            
            if not self.pull_file(qnx_result_path_alt, local_result_path_alt):
                error_msg = f"拉取结果文件 {result_file_name_alt} 也失败"
                print(error_msg)
                self.increment_retry_count(material_path_name, file_name)
                self.update_task_status(material_path_name, file_name, STATUS_FAILED, error_msg)
            else:
                local_result_path = local_result_path_alt
                result_pulled = True
        else:
            result_pulled = True
            
        # 无论文件拉取是否成功，都清理QNX系统中的文件
        self.clean_qnx_files(file_name)
        
        if result_pulled:
            print(f"文件 {material_path_name}/{file_name} 处理完成，结果保存到 {local_result_path}")
            self.update_task_status(material_path_name, file_name, STATUS_SUCCESS)
            return True
        
        return False
    
    def process_all_files(self):
        """处理所有素材文件夹中的所有文件"""
        print(f"开始处理 {len(self.material_folders)} 个素材文件夹")
        
        # 收集所有素材文件夹中的文件
        all_files = []
        for material_folder in self.material_folders:
            print(f"扫描素材文件夹: {material_folder}")
            
            # 获取素材文件夹中的所有文件
            files = [f for f in material_folder.iterdir() if f.is_file()]
            
            if not files:
                print(f"素材文件夹 {material_folder} 中没有文件！")
                continue
            
            # 添加到总文件列表
            for file_path in files:
                all_files.append((material_folder, file_path))
        
        if not all_files:
            print("所有素材文件夹中都没有文件！")
            return
        
        # 初始化CSV文件
        self.init_csv_file(all_files)
        
        # 获取需要处理的文件列表
        files_to_process = []
        for material_folder, file_path in all_files:
            if self.should_process_file(material_folder.name, file_path.name):
                files_to_process.append((material_folder, file_path))
        
        print(f"找到 {len(files_to_process)} 个文件待处理")
        
        # 用于跟踪本次处理的任务
        processed_tasks = set()
        
        # 处理文件，带进度显示
        if TQDM_AVAILABLE:
            for material_folder, file_path in tqdm(files_to_process, desc="处理进度"):
                task_key = f"{material_folder.name}/{file_path.name}"
                processed_tasks.add(task_key)
                self.process_file(material_folder, file_path)
        else:
            total = len(files_to_process)
            for i, (material_folder, file_path) in enumerate(files_to_process):
                task_key = f"{material_folder.name}/{file_path.name}"
                processed_tasks.add(task_key)
                print(f"[{i+1}/{total}] ", end="")
                self.process_file(material_folder, file_path)
        
        print("\n" + "="*50)
        
        # 打印处理结果统计
        if self.continue_from_csv:
            # 只统计本次处理的任务
            success_count = sum(1 for key, task in self.tasks.items() if key in processed_tasks and task['status'] == STATUS_SUCCESS)
            failed_count = sum(1 for key, task in self.tasks.items() if key in processed_tasks and task['status'] == STATUS_FAILED)
            skipped_count = sum(1 for key, task in self.tasks.items() if key in processed_tasks and task['status'] == STATUS_SKIPPED)
            
            processed_count = len(files_to_process)
            print(f"处理完成！总计: {len(self.tasks)}个文件, 本次处理: {processed_count}个文件, 成功: {success_count}, 失败: {failed_count}, 跳过: {skipped_count}")
        else:
            # 统计所有任务
            success_count = sum(1 for task in self.tasks.values() if task['status'] == STATUS_SUCCESS)
            failed_count = sum(1 for task in self.tasks.values() if task['status'] == STATUS_FAILED)
            skipped_count = sum(1 for task in self.tasks.values() if task['status'] == STATUS_SKIPPED)
            
            print(f"处理完成！总计: {len(self.tasks)}个文件, 成功: {success_count}, 失败: {failed_count}, 跳过: {skipped_count}")
            
        print(f"结果文件保存在: {self.output_folder}")
        print(f"任务状态记录在: {self.csv_file}")

def main():
    # 在代码中直接设置参数
    # ====================== 用户可配置参数 ======================
    # 素材文件夹路径（支持多个文件夹）
    material_folders = [
        "/Users/only/ivqa/TestData/Zone_AI_Track/tmp_lc",
        # 可以添加更多素材文件夹路径
        "/Users/only/ivqa/TestData/Zone_AI_Track/tmp",
        # "/Users/only/ivqa/TestData/Zone_AI_Track/tmp3",
    ]
    
    # 输出文件夹路径
    output_folder = "/Users/only/Desktop/工作文档/DH/path_code/gitee_code/QNX_autotest/output"
    
    # 状态文件夹路径（存放CSV文件）
    status_folder = Path(__file__).parent / "status"
    
    # QNX系统上的推送路径
    push_path = "/mnt/vendor/nfs/fota/cjm/0526/tmp"
    
    # QNX系统上的工作目录
    qnx_work_dir = "/fota/cjm/0526"
    
    # QNX系统上要执行的算法命令
    algorithm_cmd = "./trackdemo --model ./ --input ./tmp/ --testcase 1 --width 1280 --height 960"
    
    # CSV文件路径(None表示自动创建)
    # csv_file = None
    csv_file = "/Users/only/Desktop/工作文档/DH/path_code/gitee_code/QNX_autotest/status/task_status_20250528_133511的副本.csv"

    # 是否从CSV文件继续执行未完成的任务
    continue_from_csv = True
    
    # 失败重试的最大次数
    max_retries = 1
    
    # 指定设备序列号（为空则使用默认设备）
    device_serial = "192.168.43.1:5555"
    # ===========================================================
    
    # 创建QNXAutoTest实例并运行
    auto_test = QNXAutoTest(
        material_folders, 
        output_folder, 
        push_path, 
        qnx_work_dir, 
        algorithm_cmd,
        csv_file=csv_file,
        status_folder=status_folder,
        max_retries=max_retries,
        continue_from_csv=continue_from_csv,
        device_serial=device_serial
    )
    
    auto_test.process_all_files()

if __name__ == "__main__":
    main() 